#!/bin/bash

ecut_energys=(100 90 110 120 130) # the unit is Rydberg and 1 Rydberg = 13.6 eV

cpu_cores=(6 10 16 20 28 32 40)

KPTs=(1 2  3 4 5 6 7 8 9 )




test_var=KPT

echo "start to test ${test_var} with energy and time"
cat > ${test_var}_cores_energy_time << EOF
${test_var}    energy/atom    Time
EOF


for KPT in ${KPTs[@]} 
    do 
        mkdir  "${test_var}_ecut_${ecut_energys}_cpu_core${cpu_cores}_KPT_${KPT}"
        cd ${test_var}_ecut_${ecut_energys}_cpu_core${cpu_cores}_KPT_${KPT}

        echo "start to write the INPUT of ${test_var} for  ecut ${ecut_energys}_cpu_core${cpu_cores} KPT of ${KPT} in path $PWD"
        cat > INPUT << EOF
INPUT_PARAMETERS
#Parameters (1.General)
calculation         scf
#Parameters (2.Iteration)
ecutwfc             ${ecut_energys}
scf_nmax            100
symmetry            0
kspacing            0.08
#Parameters (3.Basis)
basis_type          pw
#Parameters (4.Smearing)
smearing_method     gauss
smearing_sigma      0.002
#Parameters (5.Mixing)
mixing_type         pulay
mixing_beta         0.7
cal_stress          1
cal_force           1
EOF
        echo "start to write the KPT of ecut ${ecut_energys}_cpu_core${cpu_cores} in path $PWD"
        cat > KPT << EOF
K_POINTS
0
Gamma
${KPT} ${KPT} ${KPT} 0 0 0
EOF

            cp /home/zxg/be-cu_dp_course/confs/Be/STRU/Be16_hcp.STRU ./STRU
            cp /home/zxg/be-cu_dp_course/potential/Be_ONCV_PBE-1.0.upf ./Be_ONCV_PBE-1.0.upf

            echo "start to run abacus in path $PWD"
            OMP_NUM_THREADS=1 mpirun -n ${cpu_cores} abacus | tee out.log
            
            Time=`grep "TOTAL  Time" out.log  | awk '{printf int($4)}'`
            energy_atom=`grep "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log  | awk -v stru="16" '{printf "%.8f\n",$2/stru}'`
            echo "${test_var}    $energy_atom    $Time" >> ../ecut_energy_time
        cd ..
done

echo "start to write the python script in path $PWD"

cat > ${test_var}_energy_time.py << EOF
import matplotlib.pyplot as plt

with open('${test_var}_energy_time', 'r') as f:
    data = f.readlines()
${test_var} = []
energy_atom = []
Time = []

for line in data[1:]:
    line = line.strip().split()
    ${test_var}.append(str(line[0]))
    energy_atom.append(float(line[1]))
    Time.append(int(line[2]))

fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True)
ax1.plot(${test_var}, energy_atom, marker='o')
ax1.set_ylabel('energy/atom')
ax1.set_title('Plot of ${test_var} vs energy/atom')
ax1.set_ylim(min(energy_atom), max(energy_atom))

ax2.plot(${test_var}, Time, marker='o')
ax2.set_xlabel('${test_var}')
ax2.set_ylabel('Time')
ax2.set_title('Plot of ${test_var} vs Time')
ax2.set_ylim(min(Time), max(Time))
plt.savefig('${test_var}_energy_time.png')

EOF

python ${test_var}_energy_time.py

echo "All done"

